import matplotlib.pyplot as plt
import argparse
from collections import defaultdict

# 读取并解析统计文件数据的函数
def read_statistic_file(file_path):
    data = defaultdict(lambda: defaultdict(dict))
    current_section = None
    current_region = None

    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if "Statistics" in line:
                current_section = line.split(" Statistics:")[0].strip()
                print(f"Current section found: {current_section}")
            elif "prob <= 0.5" in line:
                current_region = "prob_le_0.5"
                print(f"Current region found: {current_region}")
            elif "0.5 < prob <= 0.6" in line:
                current_region = "prob_0.5_0.6"
                print(f"Current region found: {current_region}")
            elif line.startswith("Number of records") or line == "":
                continue
            elif line.startswith("Sequences and their counts"):
                continue
            elif current_section and current_region:
                parts = line.split(":")
                if len(parts) == 2:
                    kmer = parts[0].strip()
                    count = int(parts[1].strip())
                    if kmer:
                        data[current_section][current_region][kmer] = count
                        print(f"Added kmer: {kmer} with count: {count} to {current_section} ({current_region})")

    # 打印数据结构检查
    for section, regions in data.items():
        print(f"\nSection: {section}")
        for region, kmers in regions.items():
            print(f"  Region: {region}, Number of kmers: {len(kmers)}")
    
    return data

# 获取前10个kmer及其占比
def get_top_kmers(kmer_counts, top_n=10):
    total_count = sum(kmer_counts.values())
    if total_count == 0:
        return []
    top_kmers = sorted(kmer_counts.items(), key=lambda x: x[1], reverse=True)[:top_n]
    return [(kmer, count / total_count) for kmer, count in top_kmers]

# 绘制前10个kmer的占比图
def plot_top_kmers(data, sample_label, region_key):
    top_kmers = get_top_kmers(data[sample_label][region_key])
    if not top_kmers:
        print(f"No data available for {sample_label} ({region_key}). Skipping plot.")
        return

    kmers, proportions = zip(*top_kmers)

    plt.figure(figsize=(10, 6))
    plt.bar(kmers, proportions, color='skyblue')
    plt.xlabel('K-mer')
    plt.ylabel('Proportion')
    plt.title(f'Top 10 K-mers in {sample_label} ({region_key})')
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.show()

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="读取统计文件并绘制每个区域top10的kmer占比图")
    parser.add_argument("file_path", help="输入的统计文件路径")
    args = parser.parse_args()

    # 读取统计数据
    statistic_data = read_statistic_file(args.file_path)

    # 打印整个数据结构检查
    print("\nFinal data structure:")
    for section, regions in statistic_data.items():
        for region, kmers in regions.items():
            print(f"Section: {section}, Region: {region}, Kmers: {len(kmers)}")

    # 绘制图表，针对每个区域和样本
    for sample_label in statistic_data:
        for region_key in statistic_data[sample_label]:
            print(f"Plotting data for {sample_label} ({region_key})")
            plot_top_kmers(statistic_data, sample_label, region_key)

if __name__ == "__main__":
    main()
